import cv2
import os
import sys
import numpy as np

def create_mtcnn_detector():
    """
    创建MTCNN人脸检测器，使用facenet-pytorch实现
    
    Returns:
        MTCNN检测器对象或None（如果无法创建）
    """
    try:
        # 使用facenet-pytorch实现的MTCNN
        from facenet_pytorch import MTCNN
        # 使用CPU模式，keep_all=True表示检测所有 faces
        mtcnn = MTCNN(keep_all=True, device='cpu')
        print("使用facenet-pytorch的MTCNN实现")
        return mtcnn
    except ImportError:
        print("错误: 未找到facenet-pytorch库，请安装:")
        print("  pip install facenet-pytorch")
        return None

def capture_faces(name, num_images=50, camera_id=0, output_dir="known_faces"):
    """
    使用MTCNN捕获人脸图像用于训练人脸识别模型
    
    Args:
        name (str): 人员姓名，将用作保存图像文件名的前缀
        num_images (int): 要捕获的图像数量，默认为50张
        camera_id (int): 摄像头设备ID，默认为0（第一个摄像头）
        output_dir (str): 图像保存的目录路径，默认为"known_faces"
    """
    # 创建MTCNN检测器
    face_detector = create_mtcnn_detector()
    if face_detector is None:
        return
    
    # 确保name参数使用正确的编码，特别是处理中文字符时
    if sys.platform.startswith('win'):
        # 在Windows系统上确保正确处理中文字符
        try:
            # 尝试用系统默认编码处理
            name = name.encode('gbk').decode('gbk')
        except UnicodeError:
            try:
                # 如果失败，尝试用UTF-8编码
                name = name.encode('utf-8').decode('utf-8')
            except UnicodeError:
                # 如果还失败，使用ASCII编码安全的字符串
                name = name.encode('ascii', 'ignore').decode('ascii')
    
    print(f"开始捕获人脸图像: {name}")
    print(f"目标数量: {num_images} 张")
    print(f"输出目录: {output_dir}")
    print("=" * 50)
    
    # 创建输出目录，如果目录不存在则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建目录: {output_dir}")
    
    # 初始化摄像头设备
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print(f"错误: 无法打开摄像头 {camera_id}")
        # 尝试其他设备ID (1-4)，在某些系统中摄像头可能分配了不同的ID
        for i in range(1, 5):
            print(f"尝试摄像头 {i}...")
            cap = cv2.VideoCapture(i)
            if cap.isOpened():
                print(f"成功打开摄像头 {i}")
                camera_id = i
                break
        if not cap.isOpened():
            print("无法打开任何摄像头")
            return
    
    # 设置摄像头分辨率，提高图像质量
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    # 初始化计数器
    # captured_count: 已成功捕获并保存的图像数量
    # frame_count: 处理的视频帧总数，用于控制自动捕获频率
    captured_count = 0
    frame_count = 0
    
    print("\n操作说明:")
    print("- 确保人脸正对摄像头")
    print("- 保持不同角度和表情")
    print("- 确保光线充足且均匀")
    print("- 按 'q' 或 'ESC' 键退出")
    print("- 按 'c' 键捕获当前帧中的人脸")
    print("- 自动捕获模式下会自动保存检测到的人脸")
    print("\n准备开始...")
    
    # 等待2秒让用户准备
    cv2.waitKey(2000)
    
    try:
        # 主循环：持续从摄像头读取帧并检测人脸
        while True:
            # 从摄像头读取一帧图像
            ret, frame = cap.read()
            if not ret:
                print("无法获取摄像头画面")
                break
            
            # 帧计数器递增
            frame_count += 1
            
            # 创建帧副本用于显示，避免在原始帧上绘制影响人脸检测
            display_frame = frame.copy()
            
            # 使用MTCNN检测人脸
            faces = detect_faces_mtcnn(face_detector, frame)
            
            # 在显示帧上绘制绿色矩形框标记检测到的人脸
            for (x, y, w, h) in faces:
                cv2.rectangle(display_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
            
            # 在显示帧上添加文字信息
            cv2.putText(display_frame, f"已捕获: {captured_count}/{num_images}", 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.putText(display_frame, f"按 'c' 捕获, 'q' 退出", 
                       (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            # 显示处理后的帧
            cv2.imshow('人脸捕获 (MTCNN)', display_frame)
            
            # 自动捕获模式（每30帧检查一次，且至少检测到一个人脸）
            # 这样可以避免保存过多重复的人脸图像
            if frame_count % 30 == 0 and len(faces) > 0 and captured_count < num_images:
                # 遍历所有检测到的人脸
                for i, (x, y, w, h) in enumerate(faces):
                    # 如果已达到目标数量则停止捕获
                    if captured_count >= num_images:
                        break
                        
                    # 提取人脸区域（裁剪出人脸部分）
                    face_img = frame[int(y):int(y+h), int(x):int(x+w)]
                    
                    # 生成文件名，格式为姓名+编号，如GUO01.jpg, GUO02.jpg
                    filename = f"{name.upper()}{captured_count+1:02d}.jpg"
                    # 确保文件名在Windows系统上正确编码
                    if sys.platform.startswith('win'):
                        try:
                            filename = filename.encode('utf-8').decode('utf-8')
                        except UnicodeError:
                            filename = filename.encode('ascii', 'ignore').decode('ascii')
                    filepath = os.path.join(output_dir, filename)
                    
                    # 保存图像到指定目录，确保路径编码正确
                    success = cv2.imwrite(filepath, face_img)
                    # 如果保存失败，尝试使用编码后的路径
                    if not success:
                        try:
                            # 在Windows系统上使用gbk编码处理中文路径
                            if sys.platform.startswith('win'):
                                encoded_filepath = filepath.encode('gbk').decode('gbk')
                                success = cv2.imwrite(encoded_filepath, face_img)
                            else:
                                # 在其他系统上使用utf-8编码
                                encoded_filepath = filepath.encode('utf-8').decode('utf-8')
                                success = cv2.imwrite(encoded_filepath, face_img)
                        except UnicodeError:
                            # 如果编码失败，使用英文文件名
                            filename = f"face{captured_count+1:02d}.jpg"
                            filepath = os.path.join(output_dir, filename)
                            success = cv2.imwrite(filepath, face_img)
                    if success:
                        captured_count += 1
                        print(f"已保存: {filename} (共{captured_count}张)")
                    else:
                        print(f"保存失败: {filename}")
            
            # 等待按键输入（1毫秒）
            key = cv2.waitKey(1) & 0xFF
            
            # 按 'q' 或 ESC 键退出程序
            if key == ord('q') or key == 27:  # 27 is ESC key
                break
            
            # 按 'c' 键手动捕获当前帧中的人脸（立即保存当前检测到的人脸）
            elif key == ord('c') and len(faces) > 0:
                # 遍历所有检测到的人脸
                for i, (x, y, w, h) in enumerate(faces):
                    # 如果已达到目标数量则停止捕获
                    if captured_count >= num_images:
                        break
                        
                    # 提取人脸区域
                    face_img = frame[int(y):int(y+h), int(x):int(x+w)]
                    
                    # 生成文件名
                    filename = f"{name.upper()}{captured_count+1:02d}.jpg"
                    # 确保文件名在Windows系统上正确编码
                    if sys.platform.startswith('win'):
                        try:
                            filename = filename.encode('utf-8').decode('utf-8')
                        except UnicodeError:
                            filename = filename.encode('ascii', 'ignore').decode('ascii')
                    filepath = os.path.join(output_dir, filename)
                    
                    # 保存图像，确保路径编码正确
                    success = cv2.imwrite(filepath, face_img)
                    # 如果保存失败，尝试使用编码后的路径
                    if not success:
                        try:
                            # 在Windows系统上使用gbk编码处理中文路径
                            if sys.platform.startswith('win'):
                                encoded_filepath = filepath.encode('gbk').decode('gbk')
                                success = cv2.imwrite(encoded_filepath, face_img)
                            else:
                                # 在其他系统上使用utf-8编码
                                encoded_filepath = filepath.encode('utf-8').decode('utf-8')
                                success = cv2.imwrite(encoded_filepath, face_img)
                        except UnicodeError:
                            # 如果编码失败，使用英文文件名
                            filename = f"face{captured_count+1:02d}.jpg"
                            filepath = os.path.join(output_dir, filename)
                            success = cv2.imwrite(filepath, face_img)
                    if success:
                        captured_count += 1
                        print(f"已保存: {filename} (共{captured_count}张)")
                    else:
                        print(f"保存失败: {filename}")
            
            # 检查是否达到目标数量
            if captured_count >= num_images:
                print(f"\n已完成! 成功捕获 {captured_count} 张人脸图像")
                break
    
    except KeyboardInterrupt:
        # 处理用户按Ctrl+C中断程序的情况
        print("\n用户中断程序")
    finally:
        # 释放摄像头资源并关闭所有OpenCV窗口
        cap.release()
        cv2.destroyAllWindows()
    
    # 输出最终统计信息
    print(f"\n捕获完成!")
    print(f"目标数量: {num_images}")
    print(f"实际捕获: {captured_count}")
    print(f"保存位置: {os.path.abspath(output_dir)}")

def detect_faces_mtcnn(face_detector, image):
    """
    使用MTCNN检测图像中的人脸
    
    Args:
        face_detector: MTCNN人脸检测器对象
        image: 要检测的图像
        
    Returns:
        list: 人脸边界框列表，每个元素为(x, y, w, h)格式
    """
    try:
        # 使用MTCNN检测人脸
        # 返回boxes, probs
        boxes, probs = face_detector.detect(image)
        
        # 如果没有检测到人脸
        if boxes is None or len(boxes) == 0:
            return []
            
        # 确保boxes是numpy数组
        if not isinstance(boxes, np.ndarray):
            boxes = np.array(boxes)
            
        # 如果是三维数组，取第一个元素
        if boxes.ndim == 3:
            boxes = boxes[0]
            
        # 转换边界框格式为(x, y, w, h)
        faces = []
        for box in boxes:
            if len(box) == 4:
                # 如果是(x1, y1, x2, y2)格式，转换为(x, y, w, h)
                if box[2] > box[0] and box[3] > box[1]:  # 确保坐标有效
                    x, y, x2, y2 = box
                    w, h = x2 - x, y2 - y
                    faces.append((int(x), int(y), int(w), int(h)))
                else:
                    # 可能已经是(x, y, w, h)格式
                    faces.append((int(box[0]), int(box[1]), int(box[2]), int(box[3])))
            elif len(box) == 5:
                # 有些版本可能返回5个值，前4个是边界框
                x, y, x2, y2 = box[:4]
                w, h = x2 - x, y2 - y
                faces.append((int(x), int(y), int(w), int(h)))
                
        return faces
    except Exception as e:
        print(f"人脸检测出错: {e}")
        return []

def list_captured_faces(output_dir="known_faces"):
    """
    列出指定目录中已捕获的人脸图像文件
    
    Args:
        output_dir (str): 要列出图像的目录路径，默认为"known_faces"
    """
    # 检查目录是否存在
    if not os.path.exists(output_dir):
        print(f"目录 {output_dir} 不存在")
        return
    
    # 获取目录中所有图像文件（jpg, png, jpeg格式）
    files = [f for f in os.listdir(output_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
    if not files:
        print(f"目录 {output_dir} 中没有图像文件")
        return
    
    # 打印所有图像文件名
    print(f"\n{output_dir} 中的图像文件:")
    for file in files:
        print(f"  - {file}")
    
    # 输出文件总数
    print(f"总计: {len(files)} 个文件")

def get_user_input():
    """
    获取用户输入的姓名和捕获数量
    
    Returns:
        tuple: (姓名, 数量) 或 (None, None)（输入格式错误时）
    """
    # 提示用户输入姓名和数量
    user_input = input("请输入姓名，并指定捕获数量。示例格式：GUO+50\n")
    try:
        # 解析用户输入，按"+"分割姓名和数量
        name, count_str = user_input.split('+')
        count = int(count_str)
        return name.strip(), count
    except ValueError:
        # 处理输入格式错误的情况
        print("输入格式错误，请使用格式：姓名+数量，例如：GUO+50")
        return None, None

def main():
    """
    主函数：获取用户输入并调用捕获人脸图像功能
    """
    print("人脸图像捕获工具 (MTCNN版本)")
    print("=" * 30)
    
    # 获取用户输入的姓名和数量
    name, count = get_user_input()
    if name is None or count is None:
        return
    
    # 设置输出目录为 data/name_face 格式，如 data/GUO_face/
    output_dir = os.path.join("data", f"{name}_face")
    
    # 调用捕获人脸图像的函数
    capture_faces(name, count, 0, output_dir)

if __name__ == "__main__":
    # 程序入口点
    main()